import os
import logging
from openai import OpenAI
import base64
import asyncio
from typing import Optional, Dict, Any
import aiohttp
from datetime import datetime

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 设置OpenAI客户端的日志级别为INFO
logging.getLogger("openai").setLevel(logging.INFO)
logging.getLogger("openai._base_client").setLevel(logging.INFO)
logging.getLogger("httpcore").setLevel(logging.INFO)

# 配置
MAX_RETRIES = 3
RETRY_DELAY = 2  # 秒
MODEL_TIMEOUT = 60  # 秒
MAX_VIDEO_SIZE = 150 * 1024 * 1024  # 150MB

class VideoProcessingAgent:
    def __init__(self):
        self.client = OpenAI(
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        self._session: Optional[aiohttp.ClientSession] = None
        
    async def __aenter__(self):
        self._session = aiohttp.ClientSession()
        return self
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self._session:
            await self._session.close()
    
    def _encode_video(self, video_path: str) -> str:
        """将视频文件编码为base64"""
        try:
            with open(video_path, "rb") as video_file:
                video_data = video_file.read()
                logger.info(f"视频文件大小: {len(video_data)} bytes")
                return base64.b64encode(video_data).decode("utf-8")
        except Exception as e:
            logger.error(f"视频编码失败: {str(e)}")
            raise
    
    async def _create_completion(self, video_base64: str, text_prompt: str) -> Dict[str, Any]:
        """创建模型调用"""
        try:
            completion = await asyncio.wait_for(
                self.client.chat.completions.create(
                    model="qwen-omni-turbo",
                    messages=[
                        {
                            "role": "system",
                            "content": [{"type": "text", "text": "You are a helpful assistant."}],
                        },
                        {
                            "role": "user",
                            "content": [
                                {
                                    "type": "video_url",
                                    "video_url": {"url": f"data:;base64,{video_base64}"},
                                },
                                {"type": "text", "text": text_prompt},
                            ],
                        },
                    ],
                    modalities=["text"],
                    stream=True,
                    stream_options={"include_usage": True},
                ),
                timeout=MODEL_TIMEOUT
            )
            return completion
        except asyncio.TimeoutError:
            raise TimeoutError(f"模型调用超时（{MODEL_TIMEOUT}秒）")
        except Exception as e:
            logger.error(f"模型调用失败: {str(e)}")
            raise
    
    async def process_video(self, video_path: str, text_prompt: str = "描述这个视频的具体过程") -> Dict[str, Any]:
        """
        处理视频文件并获取AI回复
        
        Args:
            video_path: 视频文件路径
            text_prompt: 文本提示词
            
        Returns:
            dict: 包含文本回复和token使用情况的字典
        """
        retry_count = 0
        last_error = None
        
        while retry_count < MAX_RETRIES:
            try:
                # 检查文件是否存在
                if not os.path.exists(video_path):
                    raise FileNotFoundError(f"视频文件不存在: {video_path}")
                
                # 检查文件大小
                file_size = os.path.getsize(video_path)
                if file_size > MAX_VIDEO_SIZE:
                    raise ValueError(f"视频文件过大，最大支持{MAX_VIDEO_SIZE/1024/1024}MB")
                
                # 编码视频文件
                base64_video = self._encode_video(video_path)
                
                # 调用模型
                completion = await self._create_completion(base64_video, text_prompt)
                
                # 处理流式响应
                text_response = ""
                usage = None
                
                # 收集所有响应
                async for chunk in completion:
                    if chunk.choices and chunk.choices[0].delta.content:
                        text_response += chunk.choices[0].delta.content
                    elif chunk.usage:
                        usage = chunk.usage
                
                # 检查结果
                if not text_response:
                    raise ValueError("模型返回的文本为空")
                
                # 记录最终结果
                logger.info("视频处理完成")
                logger.info(f"最终响应文本: {text_response}")
                if usage:
                    logger.info(f"Token使用情况: 总token数={usage.total_tokens}")
                
                result_data = {
                    "success": True,
                    "data": {
                        "text": text_response,
                        "usage": usage
                    },
                    "error": None
                }
                logger.info(f"返回视频处理结果: {result_data}")
                return result_data
                
            except Exception as e:
                last_error = e
                retry_count += 1
                if retry_count < MAX_RETRIES:
                    logger.warning(f"处理失败，{RETRY_DELAY}秒后重试 ({retry_count}/{MAX_RETRIES}): {str(e)}")
                    await asyncio.sleep(RETRY_DELAY)
                else:
                    logger.error(f"处理失败，已达到最大重试次数: {str(e)}")
                    return {
                        "success": False,
                        "data": None,
                        "error": str(last_error)
                    }

# 创建全局实例
video_agent = VideoProcessingAgent() 